package com.bw.test3;

import com.alibaba.alink.operator.common.evaluation.TuningBinaryClassMetric;
import com.alibaba.alink.operator.common.evaluation.TuningMultiClassMetric;
import com.alibaba.alink.pipeline.PipelineModel;
import com.alibaba.alink.pipeline.classification.RandomForestClassifier;
import com.alibaba.alink.pipeline.tuning.*;
import org.apache.flink.types.Row;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.classification.RandomForestPredictBatchOp;
import com.alibaba.alink.operator.batch.classification.RandomForestTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.stream.StreamOperator;
import com.alibaba.alink.operator.stream.classification.RandomForestPredictStreamOp;
import com.alibaba.alink.operator.stream.source.MemSourceStreamOp;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

// 随机森林
public class RandomForestTrainBatchOpTest {
	@Test
	public void testRandomForestTrainBatchOp() throws Exception {
		BatchOperator.setParallelism(1);

		List <Row> df = Arrays.asList(
				Row.of(1.0, "A", 0, 0, 0),
				Row.of(2.0, "B", 1, 1, 0),
				Row.of(3.0, "C", 2, 2, 1),
				Row.of(4.0, "D", 3, 3, 1)
		);

		BatchOperator <?> batchSource
				= new MemSourceBatchOp(df, " f0 double, f1 string, f2 int, f3 int, label int");

		String[] featureCols  = new String[] {
				"f0", "f1", "f2", "f3"
		};

		String label = "label";

		RandomForestClassifier rf = new RandomForestClassifier()
				.setFeatureCols(featureCols)
				.setLabelCol(label)
				.setPredictionCol("pred")
				.setPredictionDetailCol("pred_detail");


		// 超参组合
		ParamGrid paramGrid = new ParamGrid()
				.addGrid(rf, RandomForestClassifier.SUBSAMPLING_RATIO, new Double[] {1.0, 0.99, 0.98})
				.addGrid(rf, RandomForestClassifier.NUM_TREES, new Integer[] {3, 6, 9});


		// 评估
//		BinaryClassificationTuningEvaluator tuningEvaluator = new BinaryClassificationTuningEvaluator()
//				.setLabelCol(label)
//				.setPredictionDetailCol("prediction_detail")
////				.setTuningBinaryClassMetric("AUC")
//				.setTuningBinaryClassMetric(TuningBinaryClassMetric.AUC);

		// 多分类
		MultiClassClassificationTuningEvaluator tuningEvaluator = new MultiClassClassificationTuningEvaluator()
				.setLabelCol("label")
				.setPredictionDetailCol("pred_detail")
				.setTuningMultiClassMetric(TuningMultiClassMetric.ACCURACY);

		// 网络搜索
		GridSearchCV cv = new GridSearchCV()
				.setEstimator(rf)
				.setParamGrid(paramGrid)
				.setTuningEvaluator(tuningEvaluator)
				.setNumFolds(2)
				.enableLazyPrintTrainInfo("TrainInfo");

		GridSearchCVModel model = cv.fit(batchSource);

		// 得到最优的模型
		PipelineModel bestPipelineModel = model.getBestPipelineModel();

		BatchOperator<?> transform = bestPipelineModel.transform(batchSource);


		transform.print();

	}
}